import serial
import time
from typing import Optional, Tuple


class HKVR_TG_Com:
    def __init__(self, port: str, baudrate: int = 1000000, timeout: float = 1.0):
        """初始化夹爪串口通讯对象，波特率固定为1000000bps（手册3.2通讯规格要求）{insert\_element\_2\_}"""
        self.port = port
        self.baudrate = baudrate
        self.timeout = timeout
        self.ser = None
        self.is_connected = False

    def connect(self) -> bool:
        """连接串口，按手册3.1接线定义完成硬件连接后执行此步骤{insert\_element\_3\_}"""
        try:
            self.ser = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                bytesize=serial.EIGHTBITS,
                stopbits=serial.STOPBITS_ONE,
                parity=serial.PARITY_NONE,
                timeout=self.timeout
            )
            if self.ser.is_open:
                self.is_connected = True
                print(f"✅ 成功连接串口：{self.port}")
                return True
            else:
                print(f"❌ 串口{self.port}打开失败")
                return False
        except Exception as e:
            print(f"❌ 串口连接异常：{str(e)}")
            return False

    def disconnect(self) -> None:
        """断开串口连接，避免资源占用"""
        if self.ser and self.ser.is_open:
            self.ser.close()
            self.is_connected = False
            print(f"✅ 已断开串口：{self.port}")

    def calculate_checksum(self, data: bytes) -> int:
        """按手册3.3校验和规则计算：(指定字节相加) & 0xFF{insert\_element\_4\_}"""
        if len(data) == 2:  # 复位、状态、校准指令（如01 01）
            return (data[0] + data[1]) & 0xFF
        elif len(data) == 4:  # 夹取指令（如04 03 00 F4 01）
            return (data[0] + data[1] + data[2] + data[3] + data[4]) & 0xFF
        elif len(data) == 10:  # 位置控制指令（如08 10 01 00 F4 01 EC 04 FF FF）
            return sum(data) & 0xFF
        else:
            return sum(data) & 0xFF

    def send_command(self, cmd_data: bytes) -> Optional[bytes]:
        """发送指令并接收返回数据，通讯方式为“一应一答”（手册3.3要求）{insert\_element\_5\_}"""
        if not self.is_connected:
            print("❌ 串口未连接，无法发送指令")
            return None
        try:
            checksum = self.calculate_checksum(cmd_data[2:])
            full_cmd = cmd_data + bytes([checksum])
            print(f"📤 发送指令（十六进制）：{full_cmd.hex().upper()}")
            
            self.ser.flushInput()
            self.ser.flushOutput()
            self.ser.write(full_cmd)
            time.sleep(0.05)  # 预留毫秒级响应时间（手册提及毫秒级响应速度）{insert\_element\_6\_}

            # 按指令类型定义返回数据长度（手册3.3.1-3.3.7指令格式）{insert\_element\_7\_}
            response_len_map = {
                0x52: 9,   # 获取指尖力：返回9字节
                0x03: 5,   # 夹取指令：返回5字节确认
                0x01: 5,   # 复位指令：返回5字节确认
                0x08: 8,   # 获取状态：返回8字节
                0x64: 5,   # 强制校准：返回5字节确认
                0x02: 5,   # 切换位置模式：返回5字节确认
                0x10: 5    # 位置控制：返回5字节确认
            }
            response_len = response_len_map.get(cmd_data[3], 10)
            response = self.ser.read(response_len)

            if response:
                print(f"📥 接收返回（十六进制）：{response.hex().upper()}")
                return response
            else:
                print("❌ 未收到返回数据（超时）")
                return None
        except Exception as e:
            print(f"❌ 指令收发异常：{str(e)}")
            return None

    # ------------------- 核心功能方法（对应手册3.3指令）-------------------
    def get_tip_force(self) -> Optional[Tuple[float, float, float]]:
        """获取夹爪末端三维力（手册3.3.1指令）{insert\_element\_8\_}"""
        cmd_data = bytes([0xA3, 0xB4, 0x02, 0x52, 0x00])
        response = self.send_command(cmd_data)
        if response and len(response) == 9:
            fx = int.from_bytes(response[5:7], byteorder='little', signed=True) / 100
            fy = int.from_bytes(response[7:9], byteorder='little', signed=True) / 100
            fz = int.from_bytes(response[9:11], byteorder='little', signed=True) / 100
            print(f"📊 当前三维力：Fx={fx}N, Fy={fy}N, Fz={fz}N")
            return (fx, fy, fz)
        else:
            print("❌ 获取三维力失败")
            return None

    def grip(self, speed: int = 500) -> bool:
        """夹爪夹取（手册3.3.2指令，速度范围200-1200）{insert\_element\_9\_}"""
        if not 200 <= speed <= 1200:
            print(f"❌ 速度{speed}无效，需在200-1200之间")
            return False
        speed_bytes = speed.to_bytes(2, byteorder='little')
        cmd_data = bytes([0xA3, 0xB4, 0x04, 0x03, 0x00]) + speed_bytes
        response = self.send_command(cmd_data)
        result = bool(response and len(response) == 5)
        print(f"{'✅ 夹取指令发送成功' if result else '❌ 夹取指令发送失败'}")
        return result

    def reset(self) -> bool:
        """夹爪复位（手册3.3.3指令，恢复初始位置）{insert\_element\_10\_}"""
        cmd_data = bytes([0xA3, 0xB4, 0x01, 0x01])
        response = self.send_command(cmd_data)
        result = bool(response and len(response) == 5)
        print(f"{'✅ 复位指令发送成功' if result else '❌ 复位指令发送失败'}")
        return result

    def get_status(self) -> Optional[Tuple[str, int]]:
        """获取夹爪状态（手册3.3.4-3.3.5指令，状态码：00=空闲/01=夹取中/10=判断夹稳/11=已夹稳）{insert\_element\_11\_}"""
        cmd_data = bytes([0xA3, 0xB4, 0x01, 0x08])
        response = self.send_command(cmd_data)
        if response and len(response) == 8:
            status_code = response[5]
            hardness = int.from_bytes(response[6:8], byteorder='little')
            status_map = {0x00: "空闲", 0x01: "夹取中", 0x10: "判断夹稳", 0x11: "已夹稳"}
            status_desc = status_map.get(status_code, f"未知状态(0x{status_code:02X})")
            print(f"📋 当前状态：{status_desc} | 物体硬度：{hardness}")
            return (status_desc, hardness)
        else:
            print("❌ 获取状态失败")
            return None

    def force_calibrate(self) -> bool:
        """夹爪末端强制校准（手册3.3.6指令，定期校准避免零点漂移）{insert\_element\_12\_}"""
        cmd_data = bytes([0xA3, 0xB4, 0x01, 0x64])
        response = self.send_command(cmd_data)
        result = bool(response and len(response) == 5)
        print(f"{'✅ 强制校准指令发送成功' if result else '❌ 强制校准指令发送失败'}")
        return result

    def switch_position_mode(self) -> bool:
        """切换位置控制模式（手册3.3.7指令，位置控制前需执行）{insert\_element\_13\_}"""
        cmd_data = bytes([0xA3, 0xB4, 0x01, 0x02])
        response = self.send_command(cmd_data)
        result = bool(response and len(response) == 5)
        print(f"{'✅ 位置控制模式切换成功' if result else '❌ 位置控制模式切换失败'}")
        return result

    def set_position(self, position: int, speed: int = 500) -> bool:
        """位置控制（手册3.3.7指令，位置0-2940、速度200-1200）{insert\_element\_14\_}"""
        if not 0 <= position <= 2940:
            print(f"❌ 位置{position}无效，需在0-2940之间")
            return False
        if not 200 <= speed <= 1200:
            print(f"❌ 速度{speed}无效，需在200-1200之间")
            return False
        pos_bytes = position.to_bytes(2, byteorder='little')
        speed_bytes = speed.to_bytes(2, byteorder='little')
        cmd_data = bytes([0xA3, 0xB4, 0x08, 0x10, 0x01]) + speed_bytes + pos_bytes + bytes([0xFF, 0xFF])
        response = self.send_command(cmd_data)
        result = bool(response and len(response) == 5)
        print(f"{'✅ 位置控制指令发送成功' if result else '❌ 位置控制指令发送失败'}")
        return result


def select_serial_port() -> str:
    """选择串口（Windows/COM口，Linux/dev/ttyUSB0）"""
    while True:
        print("\n📌 请选择操作系统对应的串口：")
        print("1. Windows系统（默认COM10，可修改）")
        print("2. Linux系统（默认/dev/ttyUSB0，可修改）")
        choice = input("请输入选项（1/2）：").strip()
        if choice == "1":
            port = input("请输入Windows串口（如COM10，默认COM10）：").strip() or "COM10"
            return port
        elif choice == "2":
            port = input("请输入Linux串口（如/dev/ttyUSB0，默认/dev/ttyUSB0）：").strip() or "/dev/ttyUSB0"
            return port
        else:
            print("❌ 无效选项，请重新输入")


def interactive_menu(com_obj: HKVR_TG_Com) -> None:
    """交互式功能选择菜单"""
    while True:
        print("\n" + "="*50)
        print("📋 HKVR-TG三维触觉夹爪功能测试菜单")
        print("="*50)
        print("1. 测试【获取夹爪末端三维力】（手册3.3.1）")
        print("2. 测试【夹爪夹取】（手册3.3.2，可设置速度）")
        print("3. 测试【夹爪复位】（手册3.3.3）")
        print("4. 测试【获取夹爪状态】（手册3.3.4-3.3.5）")
        print("5. 测试【夹爪末端强制校准】（手册3.3.6）")
        print("6. 测试【切换位置控制模式】（手册3.3.7）")
        print("7. 测试【位置控制】（手册3.3.7，需先执行选项6）")
        print("8. 断开串口并退出程序")
        print("="*50)
        
        choice = input("请输入测试功能的选项（1-8）：").strip()

        if choice == "1":
            print("\n🔍 开始测试【获取夹爪末端三维力】")
            com_obj.get_tip_force()

        elif choice == "2":
            print("\n🔍 开始测试【夹爪夹取】")
            speed_input = input("请输入夹取速度（200-1200，默认500）：").strip()
            speed = int(speed_input) if speed_input.isdigit() else 500
            com_obj.grip(speed=speed)
            time.sleep(2)  # 等待夹取动作完成

        elif choice == "3":
            print("\n🔍 开始测试【夹爪复位】")
            com_obj.reset()
            time.sleep(2)  # 等待复位动作完成

        elif choice == "4":
            print("\n🔍 开始测试【获取夹爪状态】")
            com_obj.get_status()

        elif choice == "5":
            print("\n🔍 开始测试【夹爪末端强制校准】")
            com_obj.force_calibrate()
            time.sleep(1)  # 等待校准完成

        elif choice == "6":
            print("\n🔍 开始测试【切换位置控制模式】")
            com_obj.switch_position_mode()

        elif choice == "7":
            print("\n🔍 开始测试【位置控制】（需先切换位置模式，即选项6）")
            pos_input = input("请输入目标位置（0-2940）：").strip()
            speed_input = input("请输入运动速度（200-1200，默认500）：").strip()
            position = int(pos_input) if pos_input.isdigit() else 0
            speed = int(speed_input) if speed_input.isdigit() else 500
            com_obj.set_position(position=position, speed=speed)
            time.sleep(2)  # 等待位置移动完成

        elif choice == "8":
            print("\n📤 正在退出程序...")
            com_obj.disconnect()
            print("✅ 程序已退出")
            break

        else:
            print("❌ 无效选项，请输入1-8之间的数字")

        # 询问是否继续测试
        continue_test = input("\n是否继续测试其他功能？（y/n，默认y）：").strip().lower() or "y"
        if continue_test != "y":
            print("\n📤 正在退出程序...")
            com_obj.disconnect()
            print("✅ 程序已退出")
            break


if __name__ == "__main__":
    print("="*60)
    print("📦 HKVR-TG三维触觉夹爪交互式测试工具（基于手册Ver.1.0.0.2）")
    print("⚠️  注意：使用前需按手册3.1接线定义完成硬件连接（黑/GND、黄/RS485A、红/VCC、绿/RS485B、灰/信号GND）{insert\_element\_15\_}")
    print("="*60)

    # 1. 选择串口
    serial_port = select_serial_port()

    # 2. 初始化并连接串口
    hkvr_com = HKVR_TG_Com(port=serial_port)
    if not hkvr_com.connect():
        print("❌ 串口连接失败，程序无法继续运行")
        exit(1)

    # 3. 进入交互式菜单
    try:
        interactive_menu(hkvr_com)
    except KeyboardInterrupt:
        print("\n\n⚠️ 用户强制中断程序")
        hkvr_com.disconnect()
        print("✅ 程序已退出")
    except Exception as e:
        print(f"\n❌ 程序运行异常：{str(e)}")
        hkvr_com.disconnect()
        print("✅ 程序已退出")